package demo.kafka.interceptor;

import org.apache.kafka.clients.producer.ProducerInterceptor;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.clients.producer.RecordMetadata;
import org.apache.kafka.common.header.internals.RecordHeader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.charset.StandardCharsets;
import java.util.Map;

public class DemoProducerInterceptor implements ProducerInterceptor<String, String> {

    private final Logger log = LoggerFactory.getLogger(DemoProducerInterceptor.class);

    private static final String TRACE_ID_HEADER_KEY = "x-trace-id";

    @Override
    public ProducerRecord<String, String> onSend(ProducerRecord<String, String> producerRecord) {

        // 1. 模拟生成一个追踪ID（在实际应用中可能从MDC或请求上下文中获取）
        String traceId = java.util.UUID.randomUUID().toString();

        // 2. 将追踪ID添加到消息Header中
        ProducerRecord<String, String> modifiedRecord = new ProducerRecord<>(
                producerRecord.topic(),
                producerRecord.partition(),
                producerRecord.key(),
                producerRecord.value(),
                producerRecord.headers()
        );
        // 添加Header
        modifiedRecord.headers().add(new RecordHeader(
                TRACE_ID_HEADER_KEY,
                traceId.getBytes(StandardCharsets.UTF_8))
        );

        // 3. (可选) 甚至可以修改消息本身或目标Topic
        // String modifiedValue = "PREFIX_" + record.value();
        // return new ProducerRecord<>(record.topic(), record.key(), modifiedValue);

        return modifiedRecord;
    }

    @Override
    public void onAcknowledgement(RecordMetadata metadata, Exception exception) {
        // 在消息被服务器确认（成功发送或发送失败）后调用
        // 注意：此方法在Producer的I/O线程中调用，故不要进行阻塞操作
        if (exception != null) {
            log.error("PRODUCER INTERCEPTOR: Failed to send message to topic {}. Error: {}",
                    metadata != null ? metadata.topic() : "unknown", exception.getMessage());
            // 这里可以上报失败指标到监控系统（如Micrometer）
        } else {
            log.info("PRODUCER INTERCEPTOR: Successfully sent message to topic {}, partition {}, offset {}",
                    metadata.topic(), metadata.partition(), metadata.offset());
            // 这里可以上报成功指标到监控系统
        }
    }

    @Override
    public void close() {
        // 拦截器关闭时调用，用于清理资源
        log.info("PRODUCER INTERCEPTOR: Closing interceptor.");
    }

    @Override
    public void configure(Map<String, ?> configs) {
        // 初始化方法，可以获取Producer的配置
        log.info("PRODUCER INTERCEPTOR: Configuring interceptor with props: {}", configs);
    }
}
